#!/usr/bin/env python

import glob, os, sys
from tqdm import trange
from astropy.table import Table
from astropy.io import fits
from scipy.stats import sigmaclip
from scipy import odr
import numpy as np
from matplotlib import pyplot as plt

def remove_overscan(imgdata, ch):
	# print('> removing 71 cols overscan ...')
    img = np.zeros((4616,1152), dtype=float)

#	for i in range(4):
#		img[0:4616, i*1152:(i+1)*1152] = imgdata[0:4616,(i*1250+27):((i+1)*1250-71)]
#		img[4616:, i*1152:(i+1)*1152] = imgdata[4784:,(i*1250+27):((i+1)*1250-71)] 
#		oc1 = np.median(imgdata[0:4616,((i+1)*1250-71):((i+1)*1250)], axis=1).astype(float) 
#		oc2 = np.median(imgdata[4784:,((i+1)*1250-71):((i+1)*1250)], axis=1).astype(float)
#		# print('oc1.shape: {}'.format(oc1.shape))
#		for j in range(4616):
#			img[j,i*1152:(i+1)*1152] = img[j,i*1152:(i+1)*1152] - oc1[j]
#			img[4616+j,i*1152:(i+1)*1152] = img[4616+j,i*1152:(i+1)*1152]- oc2[j]
#	for i in range(4,8):
#		img[0:4616, i*1152:(i+1)*1152] = imgdata[0:4616,(i*1250+71):((i+1)*1250-27)]
#		img[4616:, i*1152:(i+1)*1152] = imgdata[4784:,(i*1250+71):((i+1)*1250-27)]
#		oc1 = np.median(imgdata[0:4616,(i*1250):(i*1250+71)], axis=1).astype(float)
#		oc2 = np.median(imgdata[4784:,(i*1250):(i*1250+71)], axis=1).astype(float)
#		for j in range(4616):
#			img[j, i*1152:(i+1)*1152] = img[j, i*1152:(i+1)*1152] - oc1[j]
#			img[4616+j,i*1152:(i+1)*1152] = img[4616+j,i*1152:(i+1)*1152] - oc2[j]

    if ch < 4:
        img = imgdata[0:4616,(ch*1250+27):((ch+1)*1250-71)]
        oc1 = np.median(imgdata[0:4616,((ch+1)*1250-71):((ch+1)*1250)], axis=1).astype(float)
        for j in range(4616):
            img[j,:] = img[j,:] - oc1[j]
    elif ch < 8:
        img = imgdata[0:4616,(ch*1250+71):((ch+1)*1250-27)]
        oc1 = np.median(imgdata[0:4616,(ch*1250):(ch*1250+71)], axis=1).astype(float)
        for j in range(4616):
            img[j, :] = img[j,:] - oc1[j]
    elif ch < 12:
        img = imgdata[4784:,((15-ch)*1250+71):((16-ch)*1250-27)]
        oc2 = np.median(imgdata[4784:,((15-ch)*1250):((15-ch)*1250+71)], axis=1).astype(float)
        for j in range(4616):
            img[j,:] = img[j,:] - oc2[j]
    else:
        img = imgdata[4784:,((15-ch)*1250+27):((16-ch)*1250-71)] 
        oc2 = np.median(imgdata[4784:,((16-ch)*1250-71):((16-ch)*1250)], axis=1).astype(float)
        for j in range(4616):
            img[j,:] = img[j,:]- oc2[j]

    return img

def bindata(image, n):
    ny, nx = image.shape
    nx2 = nx // n
    x_res = nx - nx2 * n
    x1 = x_res // 2
    x2 = x_res - x1
    ny2 = ny // n
    y_res = ny - ny2 * n
    y1 = y_res // 2
    y2 = y_res - y1
    data = image[y1:ny-y2, x1:nx-x2]
    # data = data.reshape((ny2, n, nx2, n)).sum(axis=1).sum(axis=-1)
    data = data.reshape((ny2, n, nx2, n)).mean(axis=1).mean(axis=-1)
    return data

#blist = glob.glob("fits/bias/*.fits")

#bias = None
#bcnt = 0
#for i in range(len(blist)):
#    if i==0:
#        bias = fits.getdata(blist[i]).astype(float)
#    else:
#        bias += fits.getdata(blist[i]).astype(float)

#    bcnt += 1

#bias = bias / bcnt


def processPTC(fits_dir, out_dir):

    if os.path.exists(out_dir) is False:
        os.system(f'mkdir {out_dir}')

    ftable = f'{out_dir}/flat_filelist.tab'
    fdata = f'{out_dir}/ptcdata.tab'

    if os.path.exists(ftable) is False:
        flist = glob.glob(f'{fits_dir}/*ptc*.fits')
        t_exp = []
        seq = []
        for f in flist:
            name = os.path.basename(f)
            tmp = name.split('_')
#            print(tmp)
            t_exp.append(float(tmp[-4].replace('s', '')))
            seq.append(int(tmp[-2].replace('.fits', '')))
        tab = Table()
        tab['name'] = flist
        tab['t_exp'] = t_exp
        tab['seq'] = seq
        tab.sort(['t_exp', 'seq'])
        tab.write(ftable, format='ipac', overwrite=True)

#    print('#########')
#    sys.exit(0)

    clip_n = 4
    tab = Table.read(ftable, format='ipac')
    if os.path.exists(fdata) is False:
        t_exp = np.unique(tab['t_exp'])
        t_exp = t_exp[t_exp < 35]
        t = []
        chans = []
        flux_1 = []
        var_1 = []
        flux_2 = []
        var_2 = []
        flux_4 = []
        var_4 = []
        flux_8 = []
        var_8 = []
        for chan in range(16):
            for i in trange(len(t_exp), desc='channel {}'.format(chan+1)):
                index = tab['t_exp'] == t_exp[i]
                f1, f2 = tab['name'][index][0], tab['name'][index][1]
                with fits.open(f1) as img1, fits.open(f2) as img2:
                    t.append(t_exp[i])
                    chans.append(chan+1)
    #                m1 = img1[0].data.astype(float)[:4617, chan*1250+72:chan*1250+1224]
    #                m2 = img2[0].data.astype(float)[:4617, chan*1250+72:chan*1250+1224]

    #                if chan < 4:
    #                    m1 = img1[0].data.astype(float)[2:4617, chan*1250+30:chan*1250+30+1170] - bias[2:4617, chan*1250+30:chan*1250+30+1170]*0
    #                    m2 = img2[0].data.astype(float)[2:4617, chan*1250+30:chan*1250+30+1170] - bias[2:4617, chan*1250+30:chan*1250+30+1170]*0              
    #                elif chan >= 4 and chan < 8:
    #                    m1 = img1[0].data.astype(float)[2:4617, chan*1250+72:chan*1250+72+1170] - bias[2:4617, chan*1250+72:chan*1250+72+1170]*0
    #                    m2 = img2[0].data.astype(float)[2:4617, chan*1250+72:chan*1250+72+1170] - bias[2:4617, chan*1250+72:chan*1250+72+1170]*0
    #                elif chan >= 8 and chan < 12:
    #                    m1 = img1[0].data.astype(float)[4785:, (15-chan)*1250+72:(15-chan)*1250+72+1170] - bias[4785:, (15-chan)*1250+72:(15-chan)*1250+72+1170]*0
    #                    m2 = img2[0].data.astype(float)[4785:, (15-chan)*1250+72:(15-chan)*1250+72+1170] - bias[4785:, (15-chan)*1250+72:(15-chan)*1250+72+1170]*0       
    #                else:
    #                    m1 = img1[0].data.astype(float)[4785:, (15-chan)*1250+30:(15-chan)*1250+30+1170] - bias[4785:, (15-chan)*1250+30:(15-chan)*1250+30+1170]*0
    #                    m2 = img2[0].data.astype(float)[4785:, (15-chan)*1250+30:(15-chan)*1250+30+1170] - bias[4785:, (15-chan)*1250+30:(15-chan)*1250+30+1170]*0

#                    if chan < 4:
#                        m1 = img1[0].data.astype(float)[2:4617, chan*1250+30:chan*1250+30+1170] #- bias[2:4617, chan*1250+30:chan*1250+30+1170]*0
#                        m2 = img2[0].data.astype(float)[2:4617, chan*1250+30:chan*1250+30+1170] #- bias[2:4617, chan*1250+30:chan*1250+30+1170]*0              
#                    elif chan >= 4 and chan < 8:
#                        m1 = img1[0].data.astype(float)[2:4617, chan*1250+72:chan*1250+72+1170] #- bias[2:4617, chan*1250+72:chan*1250+72+1170]*0
#                        m2 = img2[0].data.astype(float)[2:4617, chan*1250+72:chan*1250+72+1170] #- bias[2:4617, chan*1250+72:chan*1250+72+1170]*0
#                    elif chan >= 8 and chan < 12:
#                        m1 = img1[0].data.astype(float)[4785:, (15-chan)*1250+72:(15-chan)*1250+72+1170] #- bias[4785:, (15-chan)*1250+72:(15-chan)*1250+72+1170]*0
#                        m2 = img2[0].data.astype(float)[4785:, (15-chan)*1250+72:(15-chan)*1250+72+1170] #- bias[4785:, (15-chan)*1250+72:(15-chan)*1250+72+1170]*0       
#                    else:
#                        m1 = img1[0].data.astype(float)[4785:, (15-chan)*1250+30:(15-chan)*1250+30+1170] #- bias[4785:, (15-chan)*1250+30:(15-chan)*1250+30+1170]*0
#                        m2 = img2[0].data.astype(float)[4785:, (15-chan)*1250+30:(15-chan)*1250+30+1170] #- bias[4785:, (15-chan)*1250+30:(15-chan)*1250+30+1170]*0

                    m1 = remove_overscan(img1[0].data, chan)
                    m2 = remove_overscan(img2[0].data, chan)

                    flux_1.append(np.median(m1 + m2) / 2)
                    var_1.append(sigmaclip(m1 - m2, clip_n, clip_n)[0].var() / 2)

                    mm1 = bindata(m1, 2)
                    mm2 = bindata(m2, 2)
                    flux_2.append(np.median(mm1 + mm2) / 2)
                    var_2.append(sigmaclip(mm1 - mm2, clip_n, clip_n)[0].var() / 2 * 4)

                    mm1 = bindata(m1, 4)
                    mm2 = bindata(m2, 4)
                    flux_4.append(np.median(mm1 + mm2) / 2)
                    var_4.append(sigmaclip(mm1 - mm2, clip_n, clip_n)[0].var() / 2 * 16)

                    mm1 = bindata(m1, 8)
                    mm2 = bindata(m2, 8)
                    flux_8.append(np.median(mm1 + mm2) / 2)
                    var_8.append(sigmaclip(mm1 - mm2, clip_n, clip_n)[0].var() / 2 * 64)

        cat = Table()
        cat['t_exp'] = t
        cat['channel'] = chans
        cat['flux_1'] = flux_1
        cat['var_1'] = var_1
        cat['flux_2'] = flux_2
        cat['var_2'] = var_2
        cat['flux_4'] = flux_4
        cat['var_4'] = var_4
        cat['flux_8'] = flux_8
        cat['var_8'] = var_8
        cat.sort(['channel', 't_exp'])
        cat.write(fdata, format='ipac', overwrite=True)


    cat = Table.read(fdata, format='ipac')
    for chan in np.unique(cat['channel']):
        index = cat['channel'] == chan
        ind = np.argmin(cat['t_exp'][index])
        cat['flux_1'][index] = cat['flux_1'][index] - cat['flux_1'][index][ind]
        cat['var_1'][index] = cat['var_1'][index] - cat['var_1'][index][ind]
        cat['flux_2'][index] = cat['flux_2'][index] - cat['flux_2'][index][ind]
        cat['var_2'][index] = cat['var_2'][index] - cat['var_2'][index][ind]
        cat['flux_4'][index] = cat['flux_4'][index] - cat['flux_4'][index][ind]
        cat['var_4'][index] = cat['var_4'][index] - cat['var_4'][index][ind]
        cat['flux_8'][index] = cat['flux_8'][index] - cat['flux_8'][index][ind]
        cat['var_8'][index] = cat['var_8'][index] - cat['var_8'][index][ind]
    index = cat['t_exp'] > 0
    cat = cat[index]

    model = odr.polynomial(1)
    plt.figure(figsize=(20, 40))
    output_gains = []
    output_chans = []
    output_bins = []
    for chan in range(16):
        index = cat['channel'] == chan+1
        bins = ['1', '2', '4', '8']
        colors = ['red', 'orange', 'blue', 'green']
        handlers = []
        txts = []
        row = chan // 4
        pos1 = row * 8 + chan - row * 4 + 1
        pos2 = pos1 + 4
        for i in range(4):
            if i != 2:
                pass
            else:
                x = cat['flux_'+bins[i]][index]
                y = cat['var_'+bins[i]][index]
                print('x = ', x)
                print('y = ', y)
                idx = x > 0
                x = x[idx]
                y = y[idx]
                fw = x[np.argmax(y)]
                ind = (x > fw * 0.1) * (x < fw * 0.8)
                data = odr.Data(x[ind], y[ind], we=1/(x[ind]**2), wd=1/(y[ind]**2))
                fitter = odr.ODR(data, model, ifixb=[0, 1], beta0=[0, 1])
                res = fitter.run()
                poly = np.poly1d(res.beta[::-1])
                nonlin = (y - poly(x)) / poly(x) * 100

                plt.subplot(8, 4, pos1)
                plt.scatter(x[~ind], y[~ind], color=colors[i], zorder=1, alpha=0.25, edgecolors='none', s=20)
                s = plt.scatter(x[ind], y[ind], color=colors[i], zorder=2, edgecolors='none', s=20)
                l, = plt.plot(x, poly(x), color=colors[i], zorder=3, ls='--')
                handlers.append(s)
                txts.append('bin {}: gain = {:.3f} e-/DN'.format(bins[i], 1/res.beta[1]))

                output_gains.append(1/res.beta[1])
                output_chans.append(chan+1)
                output_bins.append(bins[i])

                plt.subplot(8, 4, pos2)
                plt.scatter(x[~ind], nonlin[~ind], color=colors[i], zorder=2, alpha=0.25, edgecolors='none', s=20)
                plt.scatter(x[ind], nonlin[ind], color=colors[i], zorder=3, edgecolors='none', s=20)
                plt.plot(x, np.zeros(len(x)), color='grey', zorder=1)


#            x = cat['flux_'+bins[i]][index]
#            y = cat['var_'+bins[i]][index]
#            print('x = ', x)
#            print('y = ', y)
#            idx = x > 0
#            x = x[idx]
#            y = y[idx]
#            fw = x[np.argmax(y)]
#            ind = (x > fw * 0.1) * (x < fw * 0.8)
#            data = odr.Data(x[ind], y[ind], we=1/(x[ind]**2), wd=1/(y[ind]**2))
#            fitter = odr.ODR(data, model, ifixb=[0, 1], beta0=[0, 1])
#            res = fitter.run()
#            poly = np.poly1d(res.beta[::-1])
#            nonlin = (y - poly(x)) / poly(x) * 100

#            plt.subplot(8, 4, pos1)
#            plt.scatter(x[~ind], y[~ind], color=colors[i], zorder=1, alpha=0.25, edgecolors='none', s=20)
#            s = plt.scatter(x[ind], y[ind], color=colors[i], zorder=2, edgecolors='none', s=20)
#            l, = plt.plot(x, poly(x), color=colors[i], zorder=3, ls='--')
#            handlers.append(s)
#            txts.append('bin {}: gain = {:.3f} e-/DN'.format(bins[i], 1/res.beta[1]))

#            output_gains.append(1/res.beta[1])
#            output_chans.append(chan+1)
#            output_bins.append(bins[i])

#            plt.subplot(8, 4, pos2)
#            plt.scatter(x[~ind], nonlin[~ind], color=colors[i], zorder=2, alpha=0.25, edgecolors='none', s=20)
#            plt.scatter(x[ind], nonlin[ind], color=colors[i], zorder=3, edgecolors='none', s=20)
#            plt.plot(x, np.zeros(len(x)), color='grey', zorder=1)            

        plt.subplot(8, 4, pos1)
        plt.xlabel('flux (DN)')
        plt.ylabel('variance (DN^2)')
        plt.title('channel {}'.format(chan+1))
    #    plt.text(4e4,6e3,'channel {}'.format(chan+1), color='limegreen', fontsize=14)
        plt.legend(handlers, txts)
        plt.grid(True)

        plt.subplot(8, 4, pos2)
    #    plt.xlabel('flux (ADU)')
        plt.xlabel('flux (DN)')
        plt.ylabel('non-linearity (%)')
        plt.title('channel {}'.format(chan+1))
    #    plt.text(4e4,8,'channel {}'.format(chan+1), color='brown', fontsize=14)
        plt.ylim([-10, 10])
        plt.grid(True)

    plt.tight_layout()
    plt.savefig(f'{out_dir}/ptc.pdf')
    plt.close()
    tab = Table()
    tab['channel'] = output_chans
    tab['bin'] = output_bins
    tab['gain'] = output_gains
    tab.write(f'{out_dir}/gain.tab', format='ipac', overwrite=True)


if __name__ == '__main__':
    processPTC(sys.argv[1], sys.argv[2])

